﻿#include "WinProcessInfo.h"
#include <QDebug>
#include <Windows.h>
#include <Psapi.h>
#include <processthreadsapi.h>

BOOL CALLBACK EnumWindowsProc(
  _In_ HWND   hwnd,
  _In_ LPARAM lParam
)
{
    //qDebug() << "EnumWindowsProc: HWnd = " << hwnd << ", lParam = " << lParam;
    if(lParam == 0) return false;
    QStringList * pList = (QStringList *)lParam;
    TCHAR lpString[256];
    if(::GetWindowText(hwnd, lpString, 255))
    {
        pList->append(QString::fromWCharArray(lpString));
        return true;
    }
    //qDebug() << "ret = false";
    return true;
}

BOOL CALLBACK EnumWindowsProc2(
  _In_ HWND   hwnd,
  _In_ LPARAM lParam
)
{
    //qDebug() << "EnumWindowsProc: HWnd = " << hwnd << ", lParam = " << lParam;
    if(lParam == 0) return false;
    QVector<HWND> * pList = (QVector<HWND> *)lParam;
    TCHAR lpString[256];
    if(::GetWindowText(hwnd, lpString, 255))
    {
        //if(QString::fromWCharArray(lpString));
        return true;
    }
    //qDebug() << "ret = false";
    return true;
}

WinProcessInfo::WinProcessInfo()
{

}

QString WinProcessInfo::PIDtoName(unsigned long processID)
{
    TCHAR szProcessName[MAX_PATH] = TEXT("<unknown>");

    // Get a handle to the process.
    HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
                                   PROCESS_VM_READ,
                                   FALSE, processID );

    // Get the process name.

    if (NULL != hProcess )
    {
        HMODULE hMod;
        DWORD cbNeeded;

        if ( EnumProcessModules( hProcess, &hMod, sizeof(hMod), &cbNeeded) )
        {
            GetModuleBaseName( hProcess, hMod, szProcessName, sizeof(szProcessName)/sizeof(TCHAR) );
        }
    }
    CloseHandle( hProcess );
    return QString::fromWCharArray(szProcessName);
}

QStringList WinProcessInfo::listNames(bool removeUnknown)
{
    QStringList names;

    // Get the list of process identifiers.
    DWORD aProcesses[1024], cbNeeded, cProcesses;
    if( !EnumProcesses( aProcesses, sizeof(aProcesses), &cbNeeded )  )
    {
        return names;
    }
    // Calculate how many process identifiers were returned.
    cProcesses = cbNeeded / sizeof(DWORD);
    for (unsigned int i = 0; i < cProcesses; i++ )
    {
        if( aProcesses[i] != 0 )
        {
            QString n = PIDtoName(aProcesses[i]);
            if(!removeUnknown || n != "<unknown>")
            {
                names.append(n);
            }
        }
    }
    return names;
}

QVector<DWORD> WinProcessInfo::listPID()
{
    QVector<unsigned long> pids;

    // Get the list of process identifiers.
    unsigned long aProcesses[1024], cbNeeded, cProcesses;
    if( !EnumProcesses( aProcesses, sizeof(aProcesses), &cbNeeded )  )
    {
        return pids;
    }
    // Calculate how many process identifiers were returned.
    cProcesses = cbNeeded / sizeof(DWORD);
    for (unsigned int i = 0; i < cProcesses; i++ )
    {
        if( aProcesses[i] != 0 )
        {
            pids.append(aProcesses[i]);
        }
    }
    return pids;
}

QVector<unsigned long> WinProcessInfo::nameToPID(QString name)
{
    QVector<unsigned long> pids;
    // Get the list of process identifiers.
    DWORD aProcesses[1024], cbNeeded, cProcesses;
    if( !EnumProcesses( aProcesses, sizeof(aProcesses), &cbNeeded )  )
    {
        return pids;
    }
    // Calculate how many process identifiers were returned.
    cProcesses = cbNeeded / sizeof(DWORD);
    for (unsigned int i = 0; i < cProcesses; i++ )
    {
        if( aProcesses[i] != 0 )
        {
            if(name == PIDtoName(aProcesses[i]))
            {
                pids.append(aProcesses[i]);
            }
        }
    }
    return pids;
}

bool WinProcessInfo::killProcess(unsigned long pid)
{
    //Kill The Process
    HANDLE hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pid );
    if (hProcess != NULL)
    {
        TerminateProcess(hProcess, 0);
        CloseHandle(hProcess);
        return true;
    }
    return false;
}

QStringList WinProcessInfo::listWindows()
{
    QStringList list;
    ::EnumWindows(EnumWindowsProc, (LPARAM)&list);
    return list;
}

bool WinProcessInfo::findWindow(QString name)
{
    QStringList list = listWindows();
    return list.contains(name);
}

int WinProcessInfo::killProcess(QString processName)
{
    QVector<unsigned long> pids = nameToPID(processName);
    QVectorIterator<unsigned long> it(pids);
    while(it.hasNext())
    {
        killProcess(it.next());
    }
    return 1;
}
